import numpy as np
# from scipy.stats import norm # No longer needed for Gaussian PDF calculation
import matplotlib.pyplot as plt
import torch

# Determine device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")

class Gaussians1D:
    def __init__(self, means, std_devs, amplitudes, colors=None):
        self.mean = torch.as_tensor(means, dtype=torch.float32).to(device)
        self.std_dev = torch.as_tensor(std_devs, dtype=torch.float32).to(device)
        self.amplitude = torch.as_tensor(amplitudes, dtype=torch.float32).to(device)
        
        if colors is not None:
            self.color = torch.as_tensor(colors, dtype=torch.float32).to(device)
        else:
            self.color = None

        self.num_gaussians = self.mean.shape[0]
        self._sqrt_2pi = torch.sqrt(torch.tensor(2.0 * np.pi, device=device, dtype=torch.float32))

    def calculate_values(self, x_values: torch.Tensor) -> torch.Tensor:
        """
        Calculates the sum of Gaussian probability densities at given x_values.
        x_values: A 1D tensor of x coordinates, shape (N_x).
        Returns: A 1D tensor of summed densities, shape (N_x).
        """
        x_values_dev = x_values.to(device) 
        
        x_reshaped = x_values_dev.unsqueeze(0)       # Shape: (1, N_x)
        mean_reshaped = self.mean.unsqueeze(1)   # Shape: (N_g, 1)
        std_dev_reshaped = self.std_dev.unsqueeze(1) # Shape: (N_g, 1)
        amp_reshaped = self.amplitude.unsqueeze(1)# Shape: (N_g, 1)

        # PDF: (1 / (std * sqrt(2*pi))) * exp(-0.5 * ((x - mean) / std)^2)
        exponent = -0.5 * torch.pow((x_reshaped - mean_reshaped) / std_dev_reshaped, 2)
        pdf_term = torch.exp(exponent) / (std_dev_reshaped * self._sqrt_2pi)
        
        all_gauss_values = amp_reshaped * pdf_term # Shape: (N_g, N_x)
        
        total_density = torch.sum(all_gauss_values, dim=0) # Shape: (N_x)
        return total_density

# Parameters
num_gaussians = 1  # Number of gaussian distributions to generate
x_range_np = np.linspace(-10, 10, 1000)  # Domain for the distributions (NumPy for plotting)
x_range_torch = torch.from_numpy(x_range_np).float().to(device) # PyTorch tensor for calculations

# Generate random gaussians parameters
np.random.seed(42)  # For reproducibility
means_list = []
std_devs_list = []
amplitudes_list = []
colors_list = []

for i in range(num_gaussians):
    mean_val: float = np.random.uniform(-5, 5)
    std_dev_val: float = np.random.uniform(0.5, 2.0)
    amplitude_val: float = np.random.uniform(1, 3)
    color_val: np.ndarray = np.random.random(3) # Random RGB color
    
    means_list.append(mean_val)
    std_devs_list.append(std_dev_val)
    amplitudes_list.append(amplitude_val)
    colors_list.append(color_val)

# Create Gaussians1D object
all_gaussians_obj = Gaussians1D(
    means=means_list,
    std_devs=std_devs_list,
    amplitudes=amplitudes_list,
    colors=colors_list
)

# Calculate the superposition of all gaussians for plotting
superposition_torch = all_gaussians_obj.calculate_values(x_range_torch)
superposition_np = superposition_torch.cpu().numpy()

# Plot the individual gaussians and their superposition
plt.figure(figsize=(10, 6))

# Plot individual gaussians
if all_gaussians_obj.color is not None:
    colors_for_plot_cpu = all_gaussians_obj.color.cpu()
else: # Fallback if colors were not provided (though current setup always provides them)
    import matplotlib.cm as cm
    colors_for_plot_mpl = cm.get_cmap('viridis', all_gaussians_obj.num_gaussians)

for i in range(all_gaussians_obj.num_gaussians):
    mean_i = all_gaussians_obj.mean[i]
    std_dev_i = all_gaussians_obj.std_dev[i]
    amplitude_i = all_gaussians_obj.amplitude[i]
    
    # Calculate PDF for this specific Gaussian for plotting
    # x_range_torch is already on device
    y_values_torch_i = amplitude_i * (1.0 / (std_dev_i * all_gaussians_obj._sqrt_2pi)) * \
                       torch.exp(-0.5 * torch.pow((x_range_torch - mean_i) / std_dev_i, 2))
    
    current_color_cpu = colors_for_plot_cpu[i].numpy() if all_gaussians_obj.color is not None else colors_for_plot_mpl(i)
    
    plt.plot(x_range_np, y_values_torch_i.cpu().numpy(), color=current_color_cpu, 
             linestyle='--', alpha=0.7,
             label=f"Gaussian {i+1}: μ={mean_i.item():.2f}, σ={std_dev_i.item():.2f}, A={amplitude_i.item():.2f}")

# Plot the superposition
plt.plot(x_range_np, superposition_np, 'k-', linewidth=2, label='Superposition')

plt.title('Random Gaussian Distributions and Their Superposition (PyTorch)')
plt.xlabel('x')
plt.ylabel('Density') # Changed from Probability Density as amplitude is not normalized for PDF
plt.grid(True, alpha=0.3)
plt.legend()
plt.tight_layout()
plt.show()

def merge_gaussians(gaussians_list: list) -> Gaussians1D:
    """
    Merges a list of Gaussians1D objects into a single Gaussians1D object.
    """
    means = torch.cat([g.mean for g in gaussians_list], dim=0)
    std_devs = torch.cat([g.std_dev for g in gaussians_list], dim=0)
    amplitudes = torch.cat([g.amplitude for g in gaussians_list], dim=0)

    weighted_means = means * amplitudes
    total_amplitude = torch.sum(amplitudes)
    new_mean = torch.sum(weighted_means) / total_amplitude if total_amplitude > 0 else torch.tensor(0.0, device=device)
    delta_means = means - new_mean
    new_std_dev = torch.sqrt(torch.sum((std_devs ** 2 + delta_means ** 2) * amplitudes) / total_amplitude) if total_amplitude > 0 else torch.tensor(1.0, device=device)
    new_amplitude = total_amplitude
    new_colors = None  # Colors are not merged in this example, but can be handled if needed

    return Gaussians1D(
        means=[new_mean.item()],
        std_devs=[new_std_dev.item()],
        amplitudes=[new_amplitude.item()],
        colors=new_colors
    )

# Create some example Gaussians to demonstrate the merge function
def demonstrate_gaussian_merging():
    print("Demonstrating Gaussian merging...")
    
    # Create two separate Gaussian objects
    gaussian_1 = Gaussians1D(
        means=[-2.0],
        std_devs=[1.0],
        amplitudes=[2.0],
        colors=[[0.2, 0.5, 0.7]]
    )
    
    gaussian_2 = Gaussians1D(
        means=[2.0],
        std_devs=[1.5],
        amplitudes=[1.5],
        colors=[[0.8, 0.3, 0.2]]
    )
    
    # Merge them
    merged_gaussian = merge_gaussians([gaussian_1, gaussian_2])
    
    # Create x range for plotting
    x_range_np = np.linspace(-8, 8, 1000)
    x_range_torch = torch.from_numpy(x_range_np).float().to(device)
    
    # Calculate values
    values_1 = gaussian_1.calculate_values(x_range_torch).cpu().numpy()
    values_2 = gaussian_2.calculate_values(x_range_torch).cpu().numpy()
    merged_values = merged_gaussian.calculate_values(x_range_torch).cpu().numpy()
    sum_values = values_1 + values_2
    
    # Plot results
    plt.figure(figsize=(12, 7))
    
    # Plot original Gaussians
    plt.plot(x_range_np, values_1, 'b--', linewidth=2, 
                label=f"Gaussian 1: μ={gaussian_1.mean.item():.2f}, σ={gaussian_1.std_dev.item():.2f}, A={gaussian_1.amplitude.item():.2f}")
    plt.plot(x_range_np, values_2, 'r--', linewidth=2,
                label=f"Gaussian 2: μ={gaussian_2.mean.item():.2f}, σ={gaussian_2.std_dev.item():.2f}, A={gaussian_2.amplitude.item():.2f}")
    
    # Plot sum of originals
    plt.plot(x_range_np, sum_values, 'g-', linewidth=2, alpha=0.7,
                label="Sum of Original Gaussians")
    
    # Plot merged Gaussian
    plt.plot(x_range_np, merged_values, 'k-', linewidth=2,
                label=f"Merged Gaussian: μ={merged_gaussian.mean.item():.2f}, σ={merged_gaussian.std_dev.item():.2f}, A={merged_gaussian.amplitude.item():.2f}")
    
    plt.title('Comparison of Original Gaussians vs Merged Gaussian')
    plt.xlabel('x')
    plt.ylabel('Density')
    plt.grid(True, alpha=0.3)
    plt.legend()
    
    # Calculate and display the MSE between the sum and the merged Gaussian
    mse = ((sum_values - merged_values) ** 2).mean()
    plt.figtext(0.5, 0.01, f"Mean Squared Error between Sum and Merged: {mse:.6f}", 
                ha="center", fontsize=12, bbox={"facecolor":"orange", "alpha":0.2})
    
    plt.tight_layout()
    plt.show()

# Run the demonstration
demonstrate_gaussian_merging()


def evaluate_gaussians(gaussians_container: Gaussians1D, x_values_tensor: torch.Tensor) -> torch.Tensor:
    """
    Evaluates the sum of Gaussian densities at given x_values.
    """
    return gaussians_container.calculate_values(x_values_tensor)

def volumetric_scattering_simulation(
    gaussians_obj: Gaussians1D, 
    num_rays: int = 10000, 
    start_x: float = -10.0, 
    max_x: float = 10.0, 
    step_size: float = 0.002
):
    """
    Simulate volumetric scattering of light rays using PyTorch for batch processing.
    """
    x_pos = torch.full((num_rays,), start_x, dtype=torch.float32, device=device)
    random_thresholds = torch.rand(num_rays, dtype=torch.float32, device=device)
    accumulated_prob = torch.zeros(num_rays, dtype=torch.float32, device=device)
    
    scattered_x_positions = torch.full((num_rays,), float('nan'), dtype=torch.float32, device=device)
    active_rays_mask = torch.ones(num_rays, dtype=torch.bool, device=device)
    
    max_steps = int((max_x - start_x) / step_size) + 2 # Safety break, +2 for buffer

    for _ in range(max_steps):
        if not torch.any(active_rays_mask):
            break

        current_x_for_active_rays = x_pos[active_rays_mask]
        
        if current_x_for_active_rays.numel() == 0: # Should be caught by torch.any(active_rays_mask)
            break

        densities_for_active_rays = evaluate_gaussians(gaussians_obj, current_x_for_active_rays)
        scatter_prob_step_active = densities_for_active_rays * step_size
        accumulated_prob[active_rays_mask] += scatter_prob_step_active

        # Identify rays that scatter in this step among the currently active ones
        newly_scattered_this_step_mask_among_active = (accumulated_prob[active_rays_mask] >= random_thresholds[active_rays_mask])
        
        # Get the global indices of these newly scattered rays
        global_newly_scattered_mask = torch.zeros_like(active_rays_mask)
        global_newly_scattered_mask[active_rays_mask] = newly_scattered_this_step_mask_among_active
        
        # Record scattering position for these rays (only if not already recorded)
        # This condition ensures we record the *first* scatter.
        # Since active_rays_mask becomes false once scattered, this is implicitly handled.
        scattered_x_positions[global_newly_scattered_mask] = x_pos[global_newly_scattered_mask]
        
        # Deactivate these newly scattered rays
        active_rays_mask[global_newly_scattered_mask] = False

        # For rays that are STILL active (did not scatter in this step):
        # Move them forward
        x_pos[active_rays_mask] += step_size
        
        # For rays that are STILL active (did not scatter, moved forward):
        # Deactivate if they moved out of bounds
        # Create a temporary mask for still_active_rays that are now out of bounds
        still_active_indices = active_rays_mask.nonzero(as_tuple=True)[0]
        if still_active_indices.numel() > 0:
            out_of_bounds_among_still_active = (x_pos[still_active_indices] >= max_x)
            # Update the main active_rays_mask for those that went out of bounds
            active_rays_mask[still_active_indices[out_of_bounds_among_still_active]] = False
    
    final_scattering_positions_gpu = scattered_x_positions[~torch.isnan(scattered_x_positions)]
    final_scattering_positions_cpu = final_scattering_positions_gpu.cpu().numpy()
    
    # Plot the histogram of scattering positions
    plt.figure(figsize=(10, 6))
    if len(final_scattering_positions_cpu) > 0:
        plt.hist(final_scattering_positions_cpu, bins=50, alpha=0.7, color='blue', density=True, label='Simulated Scattering')
    else:
        print("No scattering events recorded.")

    # Calculate and plot surviving ray percentage
    x_positions = np.linspace(start_x, max_x, 500)
    survival_percentages = []
    
    for x_pos in x_positions:
        # Count rays that scattered beyond this position
        rays_that_scattered_beyond = (final_scattering_positions_cpu > x_pos).sum() if len(final_scattering_positions_cpu) > 0 else 0
        
        # Count rays that never scattered (NaN positions)
        rays_that_never_scattered = torch.isnan(scattered_x_positions).sum().cpu().item()
        
        # Total rays still alive at this position
        surviving_rays = rays_that_scattered_beyond + rays_that_never_scattered
        survival_percentage = (surviving_rays / num_rays) * 100
        survival_percentages.append(survival_percentage)
    
    # Add a second y-axis for the survival percentage
    ax1 = plt.gca()
    ax2 = ax1.twinx()
    ax2.plot(x_positions, survival_percentages, 'g-', linewidth=2, label='Surviving Rays (%)')
    ax2.set_ylabel('Surviving Rays (%)', color='g')
    ax2.tick_params(axis='y', labelcolor='g')
    ax2.set_ylim(0, 105)  # Give a little headroom above 100%
    
    # Add the new legend entry
    lines1, labels1 = ax1.get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()
    ax1.legend(lines1 + lines2, labels1 + labels2, loc='upper right')
        
    plt.title('Distribution of First Scattering Positions (PyTorch)')
    plt.xlabel('x position')
    plt.ylabel('Probability Density')
    
    # Plot the density function for comparison on a secondary y-axis
    ax1 = plt.gca()  # Get current axis (for histogram)
    ax3 = ax1.twinx()  # Create a third y-axis on the right side
    ax3.spines['right'].set_position(('outward', 60))  # Offset the right spine of ax3 outward by 60 points
    
    # Create data for density plot
    x_for_density_plot_np = np.linspace(start_x, max_x, 1000)
    x_for_density_plot_torch = torch.from_numpy(x_for_density_plot_np).float().to(device)
    
    density_values_torch = evaluate_gaussians(gaussians_obj, x_for_density_plot_torch)
    density_values_np = density_values_torch.cpu().numpy()
    
    # Normalize density for comparison (approximate integral)
    dx_density_plot = (max_x - start_x) / (len(x_for_density_plot_np) - 1) if len(x_for_density_plot_np) > 1 else 1.0
    integral_density = np.sum(density_values_np) * dx_density_plot
    if integral_density > 1e-9:
        normalized_density = density_values_np / integral_density
    else:
        normalized_density = density_values_np  # Avoid division by zero if density is zero everywhere
        
    # Plot on the third axis
    density_line = ax3.plot(x_for_density_plot_np, normalized_density, 'r-', linewidth=2, label='Normalized Density')
    ax3.set_ylabel('Normalized Density', color='r')
    ax3.tick_params(axis='y', labelcolor='r')
    
    # Add the third axis to legend
    lines1, labels1 = ax1.get_legend_handles_labels()
    lines2, labels2 = ax2.get_legend_handles_labels()
    lines3, labels3 = ax3.get_legend_handles_labels()
    ax1.legend(lines1 + lines2 + lines3, labels1 + labels2 + labels3, loc='upper right')
    
    plt.grid(True, alpha=0.3)
    plt.legend()
    plt.tight_layout()
    plt.show()
    
    return final_scattering_positions_cpu

# Run the simulation
print(f"\nStarting volumetric scattering simulation with {300000} rays...")
scattering_results = volumetric_scattering_simulation(all_gaussians_obj, num_rays=300000)
print(f"Simulation finished. Number of scattering events recorded: {len(scattering_results)}")
if len(scattering_results) > 0:
    print(f"Mean scattering position: {np.mean(scattering_results):.2f}")

# import numpy as np
# # from scipy.stats import norm # No longer needed for Gaussian PDF calculation
# import matplotlib.pyplot as plt
# import torch

# # Determine device
# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# print(f"Using device: {device}")

# class Gaussians1D:
#     def __init__(self, means, std_devs, amplitudes, colors=None):
#         self.mean = torch.as_tensor(means, dtype=torch.float32).to(device)
#         self.std_dev = torch.as_tensor(std_devs, dtype=torch.float32).to(device)
#         self.amplitude = torch.as_tensor(amplitudes, dtype=torch.float32).to(device)
        
#         if colors is not None:
#             self.color = torch.as_tensor(colors, dtype=torch.float32).to(device)
#         else:
#             self.color = None

#         self.num_gaussians = self.mean.shape[0]
#         self._sqrt_2pi = torch.sqrt(torch.tensor(2.0 * np.pi, device=device, dtype=torch.float32))

#     def calculate_values(self, x_values: torch.Tensor) -> torch.Tensor:
#         """
#         Calculates the sum of Gaussian probability densities at given x_values.
#         x_values: A 1D tensor of x coordinates, shape (N_x).
#         Returns: A 1D tensor of summed densities, shape (N_x).
#         """
#         x_values_dev = x_values.to(device) 
        
#         x_reshaped = x_values_dev.unsqueeze(0)       # Shape: (1, N_x)
#         mean_reshaped = self.mean.unsqueeze(1)   # Shape: (N_g, 1)
#         std_dev_reshaped = self.std_dev.unsqueeze(1) # Shape: (N_g, 1)
#         amp_reshaped = self.amplitude.unsqueeze(1)# Shape: (N_g, 1)

#         # PDF: (1 / (std * sqrt(2*pi))) * exp(-0.5 * ((x - mean) / std)^2)
#         exponent = -0.5 * torch.pow((x_reshaped - mean_reshaped) / std_dev_reshaped, 2)
#         pdf_term = torch.exp(exponent) / (std_dev_reshaped * self._sqrt_2pi)
        
#         all_gauss_values = amp_reshaped * pdf_term # Shape: (N_g, N_x)
        
#         total_density = torch.sum(all_gauss_values, dim=0) # Shape: (N_x)
#         return total_density

# # Parameters
# num_gaussians = 5  # Number of gaussian distributions to generate
# x_range_np = np.linspace(-10, 10, 1000)  # Domain for the distributions (NumPy for plotting)
# x_range_torch = torch.from_numpy(x_range_np).float().to(device) # PyTorch tensor for calculations

# # Generate random gaussians parameters
# np.random.seed(42)  # For reproducibility
# means_list = []
# std_devs_list = []
# amplitudes_list = []
# colors_list = []

# for i in range(num_gaussians):
#     mean_val: float = np.random.uniform(-5, 5)
#     std_dev_val: float = np.random.uniform(0.5, 2.0)
#     amplitude_val: float = np.random.uniform(0.5, 2.0)
#     color_val: np.ndarray = np.random.random(3) # Random RGB color
    
#     means_list.append(mean_val)
#     std_devs_list.append(std_dev_val)
#     amplitudes_list.append(amplitude_val)
#     colors_list.append(color_val)

# # Create Gaussians1D object
# all_gaussians_obj = Gaussians1D(
#     means=means_list,
#     std_devs=std_devs_list,
#     amplitudes=amplitudes_list,
#     colors=colors_list
# )

# # Calculate the superposition of all gaussians for plotting
# superposition_torch = all_gaussians_obj.calculate_values(x_range_torch)
# superposition_np = superposition_torch.cpu().numpy()

# # Plot the individual gaussians and their superposition
# plt.figure(figsize=(10, 6))

# # Plot individual gaussians
# if all_gaussians_obj.color is not None:
#     colors_for_plot_cpu = all_gaussians_obj.color.cpu()
# else: # Fallback if colors were not provided (though current setup always provides them)
#     import matplotlib.cm as cm
#     colors_for_plot_mpl = cm.get_cmap('viridis', all_gaussians_obj.num_gaussians)

# for i in range(all_gaussians_obj.num_gaussians):
#     mean_i = all_gaussians_obj.mean[i]
#     std_dev_i = all_gaussians_obj.std_dev[i]
#     amplitude_i = all_gaussians_obj.amplitude[i]
    
#     # Calculate PDF for this specific Gaussian for plotting
#     # x_range_torch is already on device
#     y_values_torch_i = amplitude_i * (1.0 / (std_dev_i * all_gaussians_obj._sqrt_2pi)) * \
#                        torch.exp(-0.5 * torch.pow((x_range_torch - mean_i) / std_dev_i, 2))
    
#     current_color_cpu = colors_for_plot_cpu[i].numpy() if all_gaussians_obj.color is not None else colors_for_plot_mpl(i)
    
#     plt.plot(x_range_np, y_values_torch_i.cpu().numpy(), color=current_color_cpu, 
#              linestyle='--', alpha=0.7,
#              label=f"Gaussian {i+1}: μ={mean_i.item():.2f}, σ={std_dev_i.item():.2f}, A={amplitude_i.item():.2f}")

# # Plot the superposition
# plt.plot(x_range_np, superposition_np, 'k-', linewidth=2, label='Superposition')

# plt.title('Random Gaussian Distributions and Their Superposition (PyTorch)')
# plt.xlabel('x')
# plt.ylabel('Density') # Changed from Probability Density as amplitude is not normalized for PDF
# plt.grid(True, alpha=0.3)
# plt.legend()
# plt.tight_layout()
# plt.show()


# def evaluate_gaussians(gaussians_container: Gaussians1D, x_values_tensor: torch.Tensor) -> torch.Tensor:
#     """
#     Evaluates the sum of Gaussian densities at given x_values.
#     """
#     return gaussians_container.calculate_values(x_values_tensor)

# def volumetric_scattering_simulation(
#     gaussians_obj: Gaussians1D, 
#     num_rays: int = 10000, 
#     start_x: float = -10.0, 
#     max_x: float = 10.0, 
#     step_size: float = 0.01
# ):
#     """
#     Simulate volumetric scattering of light rays using PyTorch for batch processing.
#     """
#     x_pos = torch.full((num_rays,), start_x, dtype=torch.float32, device=device)
#     random_thresholds = torch.rand(num_rays, dtype=torch.float32, device=device)
#     accumulated_prob = torch.zeros(num_rays, dtype=torch.float32, device=device)
    
#     scattered_x_positions = torch.full((num_rays,), float('nan'), dtype=torch.float32, device=device)
#     active_rays_mask = torch.ones(num_rays, dtype=torch.bool, device=device)
    
#     max_steps = int((max_x - start_x) / step_size) + 2 # Safety break, +2 for buffer

#     for _ in range(max_steps):
#         if not torch.any(active_rays_mask):
#             break

#         current_x_for_active_rays = x_pos[active_rays_mask]
        
#         if current_x_for_active_rays.numel() == 0: # Should be caught by torch.any(active_rays_mask)
#             break

#         densities_for_active_rays = evaluate_gaussians(gaussians_obj, current_x_for_active_rays)
#         scatter_prob_step_active = densities_for_active_rays * step_size
#         accumulated_prob[active_rays_mask] += scatter_prob_step_active

#         # Identify rays that scatter in this step among the currently active ones
#         newly_scattered_this_step_mask_among_active = (accumulated_prob[active_rays_mask] >= random_thresholds[active_rays_mask])
        
#         # Get the global indices of these newly scattered rays
#         global_newly_scattered_mask = torch.zeros_like(active_rays_mask)
#         global_newly_scattered_mask[active_rays_mask] = newly_scattered_this_step_mask_among_active
        
#         # Record scattering position for these rays (only if not already recorded)
#         # This condition ensures we record the *first* scatter.
#         # Since active_rays_mask becomes false once scattered, this is implicitly handled.
#         scattered_x_positions[global_newly_scattered_mask] = x_pos[global_newly_scattered_mask]
        
#         # Deactivate these newly scattered rays
#         active_rays_mask[global_newly_scattered_mask] = False

#         # For rays that are STILL active (did not scatter in this step):
#         # Move them forward
#         x_pos[active_rays_mask] += step_size
        
#         # For rays that are STILL active (did not scatter, moved forward):
#         # Deactivate if they moved out of bounds
#         # Create a temporary mask for still_active_rays that are now out of bounds
#         still_active_indices = active_rays_mask.nonzero(as_tuple=True)[0]
#         if still_active_indices.numel() > 0:
#             out_of_bounds_among_still_active = (x_pos[still_active_indices] >= max_x)
#             # Update the main active_rays_mask for those that went out of bounds
#             active_rays_mask[still_active_indices[out_of_bounds_among_still_active]] = False
    
#     final_scattering_positions_gpu = scattered_x_positions[~torch.isnan(scattered_x_positions)]
#     final_scattering_positions_cpu = final_scattering_positions_gpu.cpu().numpy()
    
#     # Plot the histogram of scattering positions
#     plt.figure(figsize=(10, 6))
#     if len(final_scattering_positions_cpu) > 0:
#         plt.hist(final_scattering_positions_cpu, bins=50, alpha=0.7, color='blue', density=True, label='Simulated Scattering')
#     else:
#         print("No scattering events recorded.")
        
#     plt.title('Distribution of First Scattering Positions (PyTorch)')
#     plt.xlabel('x position')
#     plt.ylabel('Probability Density')
    
#     # Plot the density function for comparison
#     x_for_density_plot_np = np.linspace(start_x, max_x, 1000)
#     x_for_density_plot_torch = torch.from_numpy(x_for_density_plot_np).float().to(device)
    
#     density_values_torch = evaluate_gaussians(gaussians_obj, x_for_density_plot_torch)
#     density_values_np = density_values_torch.cpu().numpy()
    
#     # Normalize density for comparison (approximate integral)
#     dx_density_plot = (max_x - start_x) / (len(x_for_density_plot_np) - 1) if len(x_for_density_plot_np) > 1 else 1.0
#     integral_density = np.sum(density_values_np) * dx_density_plot
#     if integral_density > 1e-9:
#         normalized_density = density_values_np / integral_density
#     else:
#         normalized_density = density_values_np # Avoid division by zero if density is zero everywhere
        
#     plt.plot(x_for_density_plot_np, normalized_density, 'r-', linewidth=2, label='Normalized Summed Density')
    
#     plt.grid(True, alpha=0.3)
#     plt.legend()
#     plt.tight_layout()
#     plt.show()
    
#     return final_scattering_positions_cpu

# # Run the simulation
# print(f"\nStarting volumetric scattering simulation with {10000} rays...")
# scattering_results = volumetric_scattering_simulation(all_gaussians_obj, num_rays=10000)
# print(f"Simulation finished. Number of scattering events recorded: {len(scattering_results)}")
# if len(scattering_results) > 0:
#     print(f"Mean scattering position: {np.mean(scattering_results):.2f}")
